"""
LangGraph Edge 条件分支示例
==========================

本例演示：
1️⃣ Edge 不只是顺序，还可以带条件。
2️⃣ 通过 `add_conditional_edges` 动态决定下一个节点。
"""

from langgraph.graph import StateGraph, START, END


# --- Step 1️⃣：定义节点 ---
def start_node(state: dict):
    print("🚀 [start_node] 收到输入：", state)
    return state  # 不修改，直接流转

def positive_node(state: dict):
    print("🟢 [positive_node] 输入大于0，state：", state)
    return {"result": "positive"}

def negative_node(state: dict):
    print("🔴 [negative_node] 输入小于0，state：", state)
    return {"result": "negative"}

def zero_node(state: dict):
    print("🟡 [zero_node] 输入等于0，state：", state)
    return {"result": "zero"}


# --- Step 2️⃣：构建状态图 ---
builder = StateGraph(dict)

builder.add_node("start", start_node)
builder.add_node("positive", positive_node)
builder.add_node("negative", negative_node)
builder.add_node("zero", zero_node)

# --- Step 3️⃣：添加条件边 ---
# add_conditional_edges() 可以根据某个判断结果跳转不同节点。
# 它接收一个 "条件函数"，返回下一个节点的 key。

def route_condition(state: dict):
    """根据输入的数值决定走哪条边"""
    num = state.get("input", 0)
    if num > 0:
        return "positive"
    elif num < 0:
        return "negative"
    else:
        return "zero"

# START → start_node
builder.add_edge(START, "start")

# start_node → (根据条件决定)
builder.add_conditional_edges("start", route_condition)

# 三个分支节点都回到 END
builder.add_edge("positive", END)
builder.add_edge("negative", END)
builder.add_edge("zero", END)

# --- Step 4️⃣：编译图 ---
graph = builder.compile()


# --- Step 5️⃣：执行不同输入测试 ---
for val in [10, -3, 0]:
    print(f"\n======================\n输入值: {val}")
    result = graph.invoke({"input": val})
    print("✅ 最终结果:", result)
